from typing import Any, List, Dict

from calc_manager.base.plugins import PluginObject
from psycopg2 import pool
from dataclasses import dataclass
import os
import numpy as np

@dataclass
class G70:
    id: int
    name: str


@dataclass
class Component:
    cid: int
    name: str
    g70_name: str


@dataclass
class Property:
    pid: int
    name: str
    mids: List[int]
    card: str
    data: Dict[str, Any]


@dataclass
class Element:
    eid: int
    element_type: str
    connectivity: List[int]
    pid: int
    cid: int
    ext: Dict[str, Any]
    refine: bool
    property: Property
    component: Component


@dataclass
class Material:
    mid: int
    name: str
    card: str
    data: Dict[str, Any]


@dataclass
class Node:
    nid: int
    x: float
    y: float
    z: float
    ext: Dict[str, Any]


class BDFPlugin(PluginObject):
    def __init__(self):
        self.db_pool = pool.SimpleConnectionPool(
            minconn=1,
            maxconn=3,
            host=os.environ.get("DB_HOST"),
            port=int(os.environ.get("DB_PORT")),
            dbname=os.environ.get("DB_DBNAME"),
            user=os.environ.get("DB_USER"),
            password=os.environ.get("DB_PASSWORD")
        )
        print("bdf plugin init......")

    def get_g70(self, model_fp: str) -> List[G70]:
        ret = []
        query_sql = f'select * from "g70_{model_fp}"'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            for data in res:
                ret.append(G70(
                    id=data[0],
                    name=data[1]
                ))
        return ret

    def get_components(self, model_fp: str, g70_name: str = None) -> List[Component]:
        ret = []
        base_query = 'SELECT * FROM "comp_{model_fp}"'
        if g70_name:
            query_sql = base_query + ' WHERE g70_name = %s'
        else:
            query_sql = base_query
        query_sql = query_sql.format(model_fp=model_fp)
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                if g70_name:
                    cur.execute(query_sql, (g70_name,))
                else:
                    cur.execute(query_sql)
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            for data in res:
                ret.append(Component(
                    cid=data[0],
                    name=data[1],
                    g70_name=data[2]
                ))
        return ret

    def get_component_by_cid(self, model_fp: str, cid: int) -> Component:
        query_sql = f'SELECT * FROM "comp_{model_fp}" where cid = {cid}'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchone()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            return Component(
                cid=res[0],
                name=res[1],
                g70_name=res[2]
            )

    def get_elements(self, model_fp: str, element_types: List[str] = None, cid: int = None, refine: bool = None) -> List[Element]:
        ret = []
        param = []
        query_sql = 'select * from "elements_{model_fp}" ele ' \
                    'left join property_{model_fp} p on p.pid = ele.pid ' \
                    'left join comp_{model_fp} c on c.cid = ele.cid where 1=1'
        if element_types:
            query_sql = query_sql + ' AND etype IN %s'
            param.append(tuple([item.upper() for item in element_types]))
        if cid:
            query_sql = query_sql + ' AND cid = %s'
            param.append(cid)
        if refine:
            query_sql = query_sql + ' AND refine = %s'
            param.append(refine)
        query_sql = query_sql.format(model_fp=model_fp)
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                if len(param) > 0:
                    cur.execute(query_sql, tuple(param))
                else:
                    cur.execute(query_sql)
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            for data in res:
                ret.append(Element(
                    eid=data[0],
                    element_type=data[1],
                    connectivity=data[2],
                    pid=data[3],
                    cid=data[4],
                    ext=data[5],
                    refine=data[6],
                    property=Property(
                        pid=data[7],
                        name=data[8],
                        mids=data[9],
                        card=data[10],
                        data=data[11]
                    ),
                    component=Component(
                        cid=data[12],
                        name=data[13],
                        g70_name=data[14]
                    )
                ))
        return ret

    def get_materials(self, model_fp) -> List[Material]:
        ret = []
        query_sql = f'select * from "material_{model_fp}"'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            for data in res:
                ret.append(Material(
                    mid=data[0],
                    name=data[1],
                    card=data[2],
                    data=data[3]
                ))
        return ret

    def get_nodes(self, model_fp: str) -> List[Node]:
        ret = []
        query_sql = f'select * from "nodes_{model_fp}"'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            for data in res:
                ret.append(Node(
                    nid=data[0],
                    x=data[1],
                    y=data[2],
                    z=data[3],
                    ext=data[4]
                ))
        return ret

    def get_properties(self, model_fp: str) -> List[Property]:
        ret = []
        query_sql = f'select * from "property_{model_fp}"'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            for data in res:
                ret.append(Property(
                    pid=data[0],
                    name=data[1],
                    mids=data[2],
                    card=data[3],
                    data=data[4]
                ))
        return ret

    def get_center(self, model_fp, eid) -> List[float]:
        query_sql = f'select * from "elements_{model_fp}" where eid = {eid}'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchone()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res is None:
            return
        connectivity = res[2]
        query_sql = f'select * from "nodes_{model_fp}" where eid IN %s'
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql, (tuple(connectivity),))
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            x_list = []
            y_list = []
            z_list = []
            for data in res:
                x_list.append(data[1])
                y_list.append(data[2])
                z_list.append(data[3])
            return [np.mean(np.array(x_list)), np.mean(np.array(y_list)), np.mean(np.array(z_list))]